from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Dict
import requests  # 注意：这里仍然使用了requests库，但不用Session
from generator.generator import Generator


FULL_API_URL = "https://api.example.com/large-model/generate"  # 替换为实际的完整URL


class LLMGenerator(Generator):
    def __init__(
            self,
            max_new_tokens: int = 100,
            split_lines: bool = True,
            temperature: float = 0,
            frequency_penalty: float = 2.0,
            presence_penalty: float = 2.0,
            num_threads: int = 10,
    ):
        self.__max_new_tokens = max_new_tokens
        self.__split_lines = split_lines
        self.__temperature = temperature
        self.__frequency_penalty = frequency_penalty
        self.__presence_penalty = presence_penalty
        self.__num_threads = num_threads

    def generate(
            self,
            texts: List[str],
            model_name: str,
    ) -> List[str]:
        if self.__num_threads == 1:
            return [self.__generate(text=text, model_name=model_name) for text in texts]
        else:
            return self.__generate_multi_thread(texts=texts, model_name=model_name)

    def __generate(
            self,
            text: str,
            model_name: str,
    ) -> str:
        params = {
            "model_name": model_name,
            "max_new_tokens": self.__max_new_tokens,
            "split_lines": self.__split_lines,
            "temperature": self.__temperature,
            "frequency_penalty": self.__frequency_penalty,
            "presence_penalty": self.__presence_penalty,
        }
        data = [text]

        response = requests.post(
            url=FULL_API_URL,  # 使用完整的URL
            params=params,
            json=data,
        )

        if response.status_code == 200:
            result = response.json()
            return result[0]
        else:
            raise ValueError(f"Generator did not work! Status code: {response.status_code}")

    def __generate_multi_thread(self, texts: List[str], model_name: str) -> List[str]:
        with ThreadPoolExecutor(max_workers=self.__num_threads) as executor:
            futures = [
                executor.submit(self.__run_generator, input_text, model_name)
                for input_text in texts
            ]
        result = [future.result() for future in as_completed(futures)]
        # 注意：这里的result已经是一个列表，每个元素都是生成的文本，不需要额外的处理
        return result

    def __run_generator(self, input_text: str, model_name: str) -> str:
        return self.__generate(text=input_text, model_name=model_name)
